from pywebio import start_server
from pywebio.input import *
from pywebio.output import *
from pywebio.pin import *
from pywebio.session import *
import cpca

def main():
    # 设置页面样式
    set_env(title="地名关键词查找", output_animation=False)
    
    # 页面标题
    put_markdown("# 地名关键词查找工具")
    
    # 输入框
    text = textarea("请输入文本", rows=10, placeholder="在此输入需要查找的文本...")
    
    # 查找逻辑
    results = []
    paragraphs = text.split('\n')
    for i, para in enumerate(paragraphs):
        if not para.strip():
            continue
        df = cpca.transform([para])
        if not df.empty:
            for _, row in df.iterrows():
                if row['省'] is not None:  # 过滤掉省份为None的结果
                    # 在原文中加粗显示地名关键字
                    highlighted_para = para
                    if row['省']:
                        highlighted_para = highlighted_para.replace(row['省'], f"**{row['省']}**")
                    if row['市']:
                        highlighted_para = highlighted_para.replace(row['市'], f"**{row['市']}**")
                    if row['区']:
                        highlighted_para = highlighted_para.replace(row['区'], f"**{row['区']}**")
                    
                    # 过滤位置字段中的None值
                    location_parts = [part for part in [row['省'], row['市'], row['区']] if part]
                    location = ''.join(location_parts)
                    
                    results.append({
                        '段落': i+1,
                        '位置': location,
                        '类型': '地名',
                        '原文': highlighted_para
                    })
    
    # 显示结果
    if results:
        # 生成可下载的CSV数据
        csv_data = "段落,位置,类型,原文\n"
        for r in results:
            csv_data += f"{r['段落']},{r['位置']},{r['类型']},{r['原文'].replace('**', '')}\n"
        
        # 显示表格和下载按钮
        put_table([
            ['段落', '位置', '类型', '原文']
        ] + [
            [r['段落'], r['位置'], r['类型'], put_markdown(r['原文'])] for r in results
        ])
        put_html('''
            <style>
                table {
                    width: 100%;
                    table-layout: fixed;
                }
                th:nth-child(1) { width: 5%; }
                th:nth-child(2) { width: 15%; }
                th:nth-child(3) { width: 10%; }
                th:nth-child(4) { width: 70%; }
            </style>
        ''')
        put_button("下载结果", onclick=lambda: put_file('results.csv', csv_data.encode('utf-8')))
    else:
        put_warning("未找到任何地名信息")

if __name__ == '__main__':
    start_server(main, port=8080, debug=True)
